{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SACC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import gymnasium as gym\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import rl_utils\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions import Normal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PolicyNetContinuous(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):\n",
    "        super(PolicyNetContinuous, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)\n",
    "        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)\n",
    "        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)\n",
    "        self.action_bound = action_bound\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        mu = self.fc_mu(x)\n",
    "        std = F.softplus(self.fc_std(x))\n",
    "        dist = Normal(mu, std)\n",
    "        normal_sample = dist.rsample()  # rsample()是重参数化采样\n",
    "        log_prob = dist.log_prob(normal_sample)\n",
    "        action = torch.tanh(normal_sample)\n",
    "        # 计算tanh_normal分布的对数概率密度\n",
    "        log_prob = log_prob - torch.log(1 - torch.tanh(action).pow(2) + 1e-7)\n",
    "        action = action * self.action_bound\n",
    "        return action, log_prob\n",
    "\n",
    "\n",
    "class QValueNetContinuous(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim):\n",
    "        super(QValueNetContinuous, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim + action_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc_out = torch.nn.Linear(hidden_dim, 1)\n",
    "\n",
    "    def forward(self, x, a):\n",
    "        cat = torch.cat([x, a], dim=1)\n",
    "        x = F.relu(self.fc1(cat))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        return self.fc_out(x)\n",
    "\n",
    "\n",
    "class SACContinuous:\n",
    "    ''' 处理连续动作的SAC算法 '''\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,\n",
    "                 actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma,\n",
    "                 device):\n",
    "        self.actor = PolicyNetContinuous(state_dim, hidden_dim, action_dim, action_bound).to(device)  # 策略网络\n",
    "        self.critic_1 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)  # 第一个Q网络\n",
    "        self.critic_2 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)  # 第二个Q网络\n",
    "        self.target_critic_1 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)  # 第一个目标Q网络\n",
    "        self.target_critic_2 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)  # 第二个目标Q网络\n",
    "        # 令目标Q网络的初始参数和Q网络一样\n",
    "        self.target_critic_1.load_state_dict(self.critic_1.state_dict())\n",
    "        self.target_critic_2.load_state_dict(self.critic_2.state_dict())\n",
    "        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n",
    "        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=critic_lr)\n",
    "        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=critic_lr)\n",
    "        # 使用alpha的log值,可以使训练结果比较稳定\n",
    "        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float)\n",
    "        self.log_alpha.requires_grad = True  #对alpha求梯度\n",
    "        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=alpha_lr)\n",
    "        self.target_entropy = target_entropy  # 目标熵的大小\n",
    "        self.gamma = gamma\n",
    "        self.tau = tau\n",
    "        self.device = device\n",
    "\n",
    "    def take_action(self, state):\n",
    "        state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
    "        action = self.actor(state)[0]\n",
    "        return [action.item()]\n",
    "\n",
    "    def calc_target(self, rewards, next_states, dones, truncated):  # 计算目标Q值\n",
    "        next_actions, log_prob = self.actor(next_states)\n",
    "        entropy = -log_prob\n",
    "        q1_value = self.target_critic_1(next_states, next_actions)\n",
    "        q2_value = self.target_critic_2(next_states, next_actions)\n",
    "        next_value = torch.min(q1_value,\n",
    "                               q2_value) + self.log_alpha.exp() * entropy\n",
    "        td_target = rewards + self.gamma * next_value * (1 - dones | truncated)\n",
    "        return td_target\n",
    "\n",
    "    def soft_update(self, net, target_net):\n",
    "        for param_target, param in zip(target_net.parameters(), net.parameters()):\n",
    "            param_target.data.copy_(param_target.data * (1.0 - self.tau) + param.data * self.tau)\n",
    "\n",
    "    def update(self, transition_dict):\n",
    "        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)\n",
    "        actions = torch.tensor(transition_dict['actions'], dtype=torch.float).view(-1, 1).to(self.device)\n",
    "        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)\n",
    "        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)\n",
    "        dones = torch.tensor(transition_dict['dones'], dtype=torch.int).view(-1, 1).to(self.device)\n",
    "        truncated = torch.tensor(transition_dict['truncated'], dtype=torch.int).view(-1, 1).to(self.device)\n",
    "        rewards = (rewards + 8.0) / 8.0  # 对倒立摆环境的奖励进行重塑\n",
    "\n",
    "        # 更新两个Q网络\n",
    "        td_target = self.calc_target(rewards, next_states, dones, truncated)\n",
    "        critic_1_loss = torch.mean(F.mse_loss(self.critic_1(states, actions), td_target.detach()))\n",
    "        critic_2_loss = torch.mean(F.mse_loss(self.critic_2(states, actions), td_target.detach()))\n",
    "        self.critic_1_optimizer.zero_grad()\n",
    "        critic_1_loss.backward()\n",
    "        self.critic_1_optimizer.step()\n",
    "        self.critic_2_optimizer.zero_grad()\n",
    "        critic_2_loss.backward()\n",
    "        self.critic_2_optimizer.step()\n",
    "\n",
    "        # 更新策略网络\n",
    "        new_actions, log_prob = self.actor(states)\n",
    "        entropy = -log_prob\n",
    "        q1_value = self.critic_1(states, new_actions)\n",
    "        q2_value = self.critic_2(states, new_actions)\n",
    "        actor_loss = torch.mean(-self.log_alpha.exp() * entropy - torch.min(q1_value, q2_value))\n",
    "        self.actor_optimizer.zero_grad()\n",
    "        actor_loss.backward()\n",
    "        self.actor_optimizer.step()\n",
    "\n",
    "        # 更新alpha值\n",
    "        alpha_loss = torch.mean((entropy - self.target_entropy).detach() * self.log_alpha.exp())\n",
    "        self.log_alpha_optimizer.zero_grad()\n",
    "        alpha_loss.backward()\n",
    "        self.log_alpha_optimizer.step()\n",
    "\n",
    "        self.soft_update(self.critic_1, self.target_critic_1)\n",
    "        self.soft_update(self.critic_2, self.target_critic_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "# 环境相关\n",
    "env_name = 'Pendulum-v1'\n",
    "env = gym.make(env_name)\n",
    "random.seed(0)\n",
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "\n",
    "# SAC算法相关\n",
    "PATH = 'checkpoints/SAC_offline_PDv1.pt'\n",
    "actor_lr = 3e-4\n",
    "critic_lr = 3e-3\n",
    "alpha_lr = 3e-4\n",
    "total_epochs = 5\n",
    "total_episodes = 100\n",
    "gamma = 0.99\n",
    "tau = 0.005  # 软更新参数\n",
    "buffer_size = 100000\n",
    "minimal_size = 1000\n",
    "batch_size = 64\n",
    "target_entropy = -env.action_space.shape[0]  # 原论文认为是动作空间的负数\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(device)\n",
    "replay_buffer = rl_utils.ReplayBuffer(buffer_size)\n",
    "\n",
    "# 神经网络相关\n",
    "state_dim = env.observation_space.shape[0]\n",
    "action_dim = env.action_space.shape[0]\n",
    "action_bound = env.action_space.high[0]  # 动作最大值\n",
    "hidden_dim = 128\n",
    "\n",
    "# 其他\n",
    "agent = SACContinuous(state_dim, hidden_dim, action_dim, action_bound, \n",
    "                      actor_lr, critic_lr, alpha_lr, target_entropy, tau, \n",
    "                      gamma, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_ckp(ckp_path):\n",
    "    if os.path.exists(ckp_path):\n",
    "        checkpoint = torch.load(ckp_path)\n",
    "        s_epoch = checkpoint['epoch']\n",
    "        s_episode = checkpoint['episode']\n",
    "        agent.actor.load_state_dict(checkpoint['actor_best_weight'])\n",
    "        agent.critic_1.load_state_dict(checkpoint['critic_1_best_weight'])\n",
    "        agent.critic_2.load_state_dict(checkpoint['critic_2_best_weight'])\n",
    "        return_list = checkpoint['return_list']\n",
    "        return s_epoch, s_episode, return_list\n",
    "    else:\n",
    "        s_epoch = 0\n",
    "        s_episode = 0\n",
    "        return_list = []\n",
    "        return s_epoch, s_episode, return_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_epoch, s_episode, return_list = read_ckp(PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                             "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "总耗时: 24分钟\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "return_list = rl_utils.train_off_policy_agent(env, agent, s_epoch, total_epochs, s_episode, total_episodes, \n",
    "                                              replay_buffer, minimal_size, batch_size, return_list, PATH, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CQL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CQL:\n",
    "    ''' CQL算法 '''\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,\n",
    "                 actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma,\n",
    "                 device, beta, num_random):\n",
    "        self.actor = PolicyNetContinuous(state_dim, hidden_dim, action_dim, action_bound).to(device)\n",
    "        self.critic_1 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.critic_2 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.target_critic_1 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.target_critic_2 = QValueNetContinuous(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.target_critic_1.load_state_dict(self.critic_1.state_dict())\n",
    "        self.target_critic_2.load_state_dict(self.critic_2.state_dict())\n",
    "        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)\n",
    "        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=critic_lr)\n",
    "        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=critic_lr)\n",
    "        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float)\n",
    "        self.log_alpha.requires_grad = True  #对alpha求梯度\n",
    "        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=alpha_lr)\n",
    "        self.target_entropy = target_entropy  # 目标熵的大小\n",
    "        self.gamma = gamma\n",
    "        self.tau = tau\n",
    "\n",
    "        self.beta = beta  # CQL损失函数中的系数\n",
    "        self.num_random = num_random  # CQL中的动作采样数\n",
    "\n",
    "    def take_action(self, state):\n",
    "        state = torch.tensor([state], dtype=torch.float).to(device)\n",
    "        action = self.actor(state)[0]\n",
    "        return [action.item()]\n",
    "\n",
    "    def soft_update(self, net, target_net):\n",
    "        for param_target, param in zip(target_net.parameters(), net.parameters()):\n",
    "            param_target.data.copy_(param_target.data * (1.0 - self.tau) + param.data * self.tau)\n",
    "\n",
    "    def update(self, transition_dict):\n",
    "        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(device)\n",
    "        actions = torch.tensor(transition_dict['actions'], dtype=torch.float).view(-1, 1).to(device)\n",
    "        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(device)\n",
    "        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(device)\n",
    "        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1, 1).to(device)\n",
    "        rewards = (rewards + 8.0) / 8.0  # 对倒立摆环境的奖励进行重塑\n",
    "\n",
    "        next_actions, log_prob = self.actor(next_states)\n",
    "        entropy = -log_prob\n",
    "        q1_value = self.target_critic_1(next_states, next_actions)\n",
    "        q2_value = self.target_critic_2(next_states, next_actions)\n",
    "        next_value = torch.min(q1_value, q2_value) + self.log_alpha.exp() * entropy\n",
    "        td_target = rewards + self.gamma * next_value * (1 - dones)\n",
    "        critic_1_loss = torch.mean(F.mse_loss(self.critic_1(states, actions), td_target.detach()))\n",
    "        critic_2_loss = torch.mean(F.mse_loss(self.critic_2(states, actions), td_target.detach()))\n",
    "\n",
    "        # 以上与SAC相同,以下Q网络更新是CQL的额外部分\n",
    "        batch_size = states.shape[0]\n",
    "        random_unif_actions = torch.rand(\n",
    "            [batch_size * self.num_random, actions.shape[-1]],\n",
    "            dtype=torch.float).uniform_(-1, 1).to(device)\n",
    "        random_unif_log_pi = np.log(0.5**next_actions.shape[-1])\n",
    "        tmp_states = states.unsqueeze(1).repeat(1, self.num_random,1).view(-1, states.shape[-1])\n",
    "        tmp_next_states = next_states.unsqueeze(1).repeat(\n",
    "            1, self.num_random, 1).view(-1, next_states.shape[-1])\n",
    "        random_curr_actions, random_curr_log_pi = self.actor(tmp_states)\n",
    "        random_next_actions, random_next_log_pi = self.actor(tmp_next_states)\n",
    "        q1_unif = self.critic_1(tmp_states, random_unif_actions).view(-1, self.num_random, 1)\n",
    "        q2_unif = self.critic_2(tmp_states, random_unif_actions).view(-1, self.num_random, 1)\n",
    "        q1_curr = self.critic_1(tmp_states, random_curr_actions).view(-1, self.num_random, 1)\n",
    "        q2_curr = self.critic_2(tmp_states, random_curr_actions).view(-1, self.num_random, 1)\n",
    "        q1_next = self.critic_1(tmp_states, random_next_actions).view(-1, self.num_random, 1)\n",
    "        q2_next = self.critic_2(tmp_states, random_next_actions).view(-1, self.num_random, 1)\n",
    "        q1_cat = torch.cat([\n",
    "            q1_unif - random_unif_log_pi,\n",
    "            q1_curr - random_curr_log_pi.detach().view(-1, self.num_random, 1),\n",
    "            q1_next - random_next_log_pi.detach().view(-1, self.num_random, 1)\n",
    "        ],\n",
    "            dim=1)\n",
    "        q2_cat = torch.cat([\n",
    "            q2_unif - random_unif_log_pi,\n",
    "            q2_curr - random_curr_log_pi.detach().view(-1, self.num_random, 1),\n",
    "            q2_next - random_next_log_pi.detach().view(-1, self.num_random, 1)\n",
    "        ],\n",
    "            dim=1)\n",
    "\n",
    "        qf1_loss_1 = torch.logsumexp(q1_cat, dim=1).mean()\n",
    "        qf2_loss_1 = torch.logsumexp(q2_cat, dim=1).mean()\n",
    "        qf1_loss_2 = self.critic_1(states, actions).mean()\n",
    "        qf2_loss_2 = self.critic_2(states, actions).mean()\n",
    "        qf1_loss = critic_1_loss + self.beta * (qf1_loss_1 - qf1_loss_2)\n",
    "        qf2_loss = critic_2_loss + self.beta * (qf2_loss_1 - qf2_loss_2)\n",
    "\n",
    "        self.critic_1_optimizer.zero_grad()\n",
    "        qf1_loss.backward(retain_graph=True)\n",
    "        self.critic_1_optimizer.step()\n",
    "        self.critic_2_optimizer.zero_grad()\n",
    "        qf2_loss.backward(retain_graph=True)\n",
    "        self.critic_2_optimizer.step()\n",
    "\n",
    "        # 更新策略网络\n",
    "        new_actions, log_prob = self.actor(states)\n",
    "        entropy = -log_prob\n",
    "        q1_value = self.critic_1(states, new_actions)\n",
    "        q2_value = self.critic_2(states, new_actions)\n",
    "        actor_loss = torch.mean(-self.log_alpha.exp() * entropy - torch.min(q1_value, q2_value))\n",
    "        self.actor_optimizer.zero_grad()\n",
    "        actor_loss.backward()\n",
    "        self.actor_optimizer.step()\n",
    "\n",
    "        # 更新alpha值\n",
    "        alpha_loss = torch.mean((entropy - self.target_entropy).detach() * self.log_alpha.exp())\n",
    "        self.log_alpha_optimizer.zero_grad()\n",
    "        alpha_loss.backward()\n",
    "        self.log_alpha_optimizer.step()\n",
    "\n",
    "        self.soft_update(self.critic_1, self.target_critic_1)\n",
    "        self.soft_update(self.critic_2, self.target_critic_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(0)\n",
    "np.random.seed(0)\n",
    "env.seed(0)\n",
    "torch.manual_seed(0)\n",
    "\n",
    "beta = 5.0\n",
    "num_random = 5\n",
    "num_epochs = 100\n",
    "num_trains_per_epoch = 500\n",
    "\n",
    "agent = CQL(state_dim, hidden_dim, action_dim, action_bound, actor_lr,\n",
    "            critic_lr, alpha_lr, target_entropy, tau, gamma, device, beta,\n",
    "            num_random)\n",
    "\n",
    "return_list = []\n",
    "for i in range(10):\n",
    "    with tqdm(total=int(num_epochs / 10), desc='Iteration %d' % i) as pbar:\n",
    "        for i_epoch in range(int(num_epochs / 10)):\n",
    "            # 此处与环境交互只是为了评估策略,最后作图用,不会用于训练\n",
    "            epoch_return = 0\n",
    "            state = env.reset()\n",
    "            done = False\n",
    "            while not done:\n",
    "                action = agent.take_action(state)\n",
    "                next_state, reward, done, _ = env.step(action)\n",
    "                state = next_state\n",
    "                epoch_return += reward\n",
    "            return_list.append(epoch_return)\n",
    "\n",
    "            for _ in range(num_trains_per_epoch):\n",
    "                b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)\n",
    "                transition_dict = {\n",
    "                    'states': b_s,\n",
    "                    'actions': b_a,\n",
    "                    'next_states': b_ns,\n",
    "                    'rewards': b_r,\n",
    "                    'dones': b_d\n",
    "                }\n",
    "                agent.update(transition_dict)\n",
    "\n",
    "            if (i_epoch + 1) % 10 == 0:\n",
    "                pbar.set_postfix({\n",
    "                    'epoch':\n",
    "                    '%d' % (num_epochs / 10 * i + i_epoch + 1),\n",
    "                    'return':\n",
    "                    '%.3f' % np.mean(return_list[-10:])\n",
    "                })\n",
    "            pbar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rl_utils.picture_return(return_list, 'SAC-CQL', env_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
